"""add expires_at to orders

Revision ID: xxx
Revises: previous_revision_id
Create Date: 2024-03-27 22:30:00

"""
from alembic import op
import sqlalchemy as sa
from datetime import datetime, timedelta
from pytz import timezone

# revision identifiers, used by Alembic.
revision = 'xxx'  # 这里会自动生成
down_revision = 'previous_revision_id'  # 需要替换为你的上一个迁移版本ID
branch_labels = None
depends_on = None

def upgrade():
    # 添加 expires_at 列
    op.add_column('orders', sa.Column('expires_at', sa.DateTime(timezone=True), nullable=True))
    
    # 更新现有订单的过期时间
    connection = op.get_bind()
    orders_table = sa.Table(
        'orders',
        sa.MetaData(),
        sa.Column('id', sa.Integer),
        sa.Column('created_at', sa.DateTime),
        sa.Column('expires_at', sa.DateTime)
    )
    
    # 获取所有订单
    orders = connection.execute(sa.select(orders_table)).fetchall()
    
    # 更新每个订单的过期时间
    beijing_tz = timezone('Asia/Shanghai')
    for order in orders:
        if order.created_at:
            created_at = order.created_at.astimezone(beijing_tz) if order.created_at.tzinfo else beijing_tz.localize(order.created_at)
            expires_at = created_at + timedelta(minutes=30)
            connection.execute(
                orders_table.update().
                where(orders_table.c.id == order.id).
                values(expires_at=expires_at)
            )
    
    # 将列设置为非空
    op.alter_column('orders', 'expires_at',
                    existing_type=sa.DateTime(timezone=True),
                    nullable=False)

def downgrade():
    op.drop_column('orders', 'expires_at') 